import pygame
import sys
import numpy as np
import time
from typing import Tuple, List, Optional

# 游戏常量定义
BOARD_SIZE = 19  # 棋盘大小
CELL_SIZE = 30   # 每个格子的大小
MARGIN = 40      # 边距
SCREEN_WIDTH = BOARD_SIZE * CELL_SIZE + MARGIN * 2
SCREEN_HEIGHT = BOARD_SIZE * CELL_SIZE + MARGIN * 2

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BOARD_COLOR = (240, 190, 120)  # 棋盘颜色
RED = (255, 0, 0)
GREEN = (0, 200, 0)

# 棋子类型定义
EMPTY = 0
PLAYER = 1  # 玩家(黑棋)
COMPUTER = 2  # 电脑(白棋)

class GomokuGame:
    def __init__(self):
        # 初始化Pygame
        pygame.init()
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("19×19五子棋")
        self.clock = pygame.time.Clock()
        
        # 初始化游戏状态
        self.board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=int)
        self.current_player = PLAYER  # 黑棋先行
        self.game_over = False
        self.winner = None
        self.last_move = None  # 记录最后一步
        
        # 加载字体
        pygame.font.init()
        self.font = pygame.font.SysFont("SimHei", 24)
        self.large_font = pygame.font.SysFont("SimHei", 48)
        
        # 游戏主循环标志
        self.running = True
        
        # 绘制初始界面
        self.draw_board()
    
    def draw_board(self):
        """绘制棋盘和当前状态"""
        # 绘制背景
        self.screen.fill(BOARD_COLOR)
        
        # 绘制网格线
        for i in range(BOARD_SIZE):
            # 横线
            pygame.draw.line(
                self.screen, BLACK, 
                (MARGIN, MARGIN + i * CELL_SIZE), 
                (MARGIN + (BOARD_SIZE - 1) * CELL_SIZE, MARGIN + i * CELL_SIZE),
                1
            )
            # 竖线
            pygame.draw.line(
                self.screen, BLACK, 
                (MARGIN + i * CELL_SIZE, MARGIN), 
                (MARGIN + i * CELL_SIZE, MARGIN + (BOARD_SIZE - 1) * CELL_SIZE),
                1
            )
        
        # 绘制天元和星位
        star_points = [(3, 3), (3, 9), (3, 15),
                      (9, 3), (9, 9), (9, 15),
                      (15, 3), (15, 9), (15, 15)]
        
        for (x, y) in star_points:
            pygame.draw.circle(
                self.screen, BLACK, 
                (MARGIN + x * CELL_SIZE, MARGIN + y * CELL_SIZE),
                5
            )
        
        # 绘制棋子
        for i in range(BOARD_SIZE):
            for j in range(BOARD_SIZE):
                if self.board[i][j] != EMPTY:
                    color = BLACK if self.board[i][j] == PLAYER else WHITE
                    pygame.draw.circle(
                        self.screen, color, 
                        (MARGIN + j * CELL_SIZE, MARGIN + i * CELL_SIZE),
                        CELL_SIZE // 2 - 2
                    )
                    # 棋子边框
                    pygame.draw.circle(
                        self.screen, BLACK, 
                        (MARGIN + j * CELL_SIZE, MARGIN + i * CELL_SIZE),
                        CELL_SIZE // 2 - 2,
                        1
                    )
        
        # 标记最后一步
        if self.last_move:
            row, col = self.last_move
            pygame.draw.circle(
                self.screen, RED, 
                (MARGIN + col * CELL_SIZE, MARGIN + row * CELL_SIZE),
                5
            )
        
        # 显示当前玩家
        if not self.game_over:
            player_text = "轮到你下棋" if self.current_player == PLAYER else "电脑思考中..."
            text_surface = self.font.render(player_text, True, BLACK)
            self.screen.blit(text_surface, (MARGIN, 10))
        else:
            # 显示游戏结果
            if self.winner == PLAYER:
                result_text = "恭喜，你赢了！"
            elif self.winner == COMPUTER:
                result_text = "很遗憾，电脑赢了！"
            else:
                result_text = "游戏结束"
            
            text_surface = self.large_font.render(result_text, True, RED)
            text_rect = text_surface.get_rect(
                center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 - 50)
            )
            self.screen.blit(text_surface, text_rect)
            
            # 显示重新开始提示
            restart_text = "按R键重新开始"
            restart_surface = self.font.render(restart_text, True, BLACK)
            restart_rect = restart_surface.get_rect(
                center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 20)
            )
            self.screen.blit(restart_surface, restart_rect)
        
        # 更新显示
        pygame.display.flip()
    
    def get_click_position(self, pos: Tuple[int, int]) -> Optional[Tuple[int, int]]:
        """根据鼠标点击位置获取棋盘坐标"""
        x, y = pos
        
        # 计算最近的交叉点
        col = round((x - MARGIN) / CELL_SIZE)
        row = round((y - MARGIN) / CELL_SIZE)
        
        # 检查是否在棋盘范围内
        if 0 <= row < BOARD_SIZE and 0 <= col < BOARD_SIZE:
            # 检查点击位置是否足够接近交叉点
            actual_x = MARGIN + col * CELL_SIZE
            actual_y = MARGIN + row * CELL_SIZE
            
            if abs(x - actual_x) <= CELL_SIZE // 2 and abs(y - actual_y) <= CELL_SIZE // 2:
                return (row, col)
        
        return None
    
    def make_move(self, row: int, col: int, player: int) -> bool:
        """在指定位置落子"""
        if self.board[row][col] == EMPTY and not self.game_over:
            self.board[row][col] = player
            self.last_move = (row, col)
            
            # 检查胜负
            if self.check_win(row, col, player):
                self.game_over = True
                self.winner = player
            
            return True
        return False
    
    def check_win(self, row: int, col: int, player: int) -> bool:
        """检查是否获胜"""
        # 检查四个方向：水平、垂直、对角线、反对角线
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        
        for dx, dy in directions:
            count = 1  # 当前位置已经有一个棋子
            
            # 正方向检查
            for i in range(1, 5):
                new_row = row + i * dx
                new_col = col + i * dy
                if (0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE and
                    self.board[new_row][new_col] == player):
                    count += 1
                else:
                    break
            
            # 反方向检查
            for i in range(1, 5):
                new_row = row - i * dx
                new_col = col - i * dy
                if (0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE and
                    self.board[new_row][new_col] == player):
                    count += 1
                else:
                    break
            
            # 五子连珠
            if count >= 5:
                return True
        
        return False
    
    def evaluate_position(self, row: int, col: int, player: int) -> int:
        """评估位置的价值"""
        score = 0
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        
        for dx, dy in directions:
            # 计算连续的棋子数和空位
            continuous = 0
            empty = 0
            blocked = 0
            
            # 正方向
            for i in range(1, 5):
                new_row = row + i * dx
                new_col = col + i * dy
                if 0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE:
                    if self.board[new_row][new_col] == player:
                        continuous += 1
                    elif self.board[new_row][new_col] == EMPTY:
                        empty += 1
                        break
                    else:
                        blocked += 1
                        break
                else:
                    blocked += 1
                    break
            
            # 反方向
            for i in range(1, 5):
                new_row = row - i * dx
                new_col = col - i * dy
                if 0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE:
                    if self.board[new_row][new_col] == player:
                        continuous += 1
                    elif self.board[new_row][new_col] == EMPTY:
                        empty += 1
                        break
                    else:
                        blocked += 1
                        break
                else:
                    blocked += 1
                    break
            
            # 评分规则
            if continuous >= 4:
                score += 100000  # 五连珠
            elif continuous == 3 and blocked == 0:
                score += 10000   # 活四
            elif continuous == 3 and blocked == 1:
                score += 1000    # 冲四
            elif continuous == 2 and blocked == 0:
                score += 1000    # 活三
            elif continuous == 2 and blocked == 1:
                score += 100     # 冲三
            elif continuous == 1 and blocked == 0:
                score += 10      # 活二
            
            # 防守评分
            opponent = PLAYER if player == COMPUTER else COMPUTER
            opponent_continuous = 0
            opponent_blocked = 0
            
            # 正方向
            for i in range(1, 5):
                new_row = row + i * dx
                new_col = col + i * dy
                if 0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE:
                    if self.board[new_row][new_col] == opponent:
                        opponent_continuous += 1
                    elif self.board[new_row][new_col] == EMPTY:
                        break
                    else:
                        opponent_blocked += 1
                        break
                else:
                    opponent_blocked += 1
                    break
            
            # 反方向
            for i in range(1, 5):
                new_row = row - i * dx
                new_col = col - i * dy
                if 0 <= new_row < BOARD_SIZE and 0 <= new_col < BOARD_SIZE:
                    if self.board[new_row][new_col] == opponent:
                        opponent_continuous += 1
                    elif self.board[new_row][new_col] == EMPTY:
                        break
                    else:
                        opponent_blocked += 1
                        break
                else:
                    opponent_blocked += 1
                    break
            
            # 防守评分
            if opponent_continuous >= 4:
                score += 90000   # 防止五连珠
            elif opponent_continuous == 3 and opponent_blocked == 0:
                score += 8000    # 防止活四
            elif opponent_continuous == 3 and opponent_blocked == 1:
                score += 800     # 防止冲四
            elif opponent_continuous == 2 and opponent_blocked == 0:
                score += 800     # 防止活三
        
        # 中心位置加成
        if abs(row - BOARD_SIZE // 2) <= 3 and abs(col - BOARD_SIZE // 2) <= 3:
            score += 5
        
        return score
    
    def get_computer_move(self) -> Tuple[int, int]:
        """获取电脑的最佳落子位置"""
        best_score = -float('inf')
        best_move = None
        
        # 收集所有可能的落子位置
        possible_moves = []
        for i in range(BOARD_SIZE):
            for j in range(BOARD_SIZE):
                if self.board[i][j] == EMPTY:
                    # 只考虑有棋子周围的空位，提高效率
                    has_neighbor = False
                    for dx in [-2, -1, 0, 1, 2]:
                        for dy in [-2, -1, 0, 1, 2]:
                            if 0 <= i+dx < BOARD_SIZE and 0 <= j+dy < BOARD_SIZE:
                                if self.board[i+dx][j+dy] != EMPTY:
                                    has_neighbor = True
                                    break
                        if has_neighbor:
                            break
                    
                    if has_neighbor or (i == BOARD_SIZE // 2 and j == BOARD_SIZE // 2):
                        possible_moves.append((i, j))
        
        # 如果棋盘为空，优先选择天元
        if not possible_moves:
            return (BOARD_SIZE // 2, BOARD_SIZE // 2)
        
        # 评估每个可能的落子位置
        for (row, col) in possible_moves:
            score = self.evaluate_position(row, col, COMPUTER)
            
            # 临时落子并评估
            self.board[row][col] = COMPUTER
            if self.check_win(row, col, COMPUTER):
                self.board[row][col] = EMPTY
                return (row, col)  # 如果能赢，直接下
            self.board[row][col] = EMPTY
            
            # 检查是否需要防守
            self.board[row][col] = PLAYER
            if self.check_win(row, col, PLAYER):
                self.board[row][col] = EMPTY
                return (row, col)  # 如果需要防守，直接下
            self.board[row][col] = EMPTY
            
            if score > best_score:
                best_score = score
                best_move = (row, col)
        
        # 如果没有找到最佳位置，随机选择一个
        if best_move is None and possible_moves:
            best_move = possible_moves[0]
        
        return best_move
    
    def reset_game(self):
        """重置游戏"""
        self.board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=int)
        self.current_player = PLAYER
        self.game_over = False
        self.winner = None
        self.last_move = None
        self.draw_board()
    
    def run(self):
        """游戏主循环"""
        while self.running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r:
                        self.reset_game()
                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    # 玩家落子
                    if not self.game_over and self.current_player == PLAYER:
                        pos = pygame.mouse.get_pos()
                        board_pos = self.get_click_position(pos)
                        if board_pos:
                            row, col = board_pos
                            if self.make_move(row, col, PLAYER):
                                self.draw_board()
                                
                                # 切换到电脑回合
                                if not self.game_over:
                                    self.current_player = COMPUTER
                                    self.draw_board()  # 显示"电脑思考中..."
                                    
                                    # 稍微延迟，让玩家看到状态变化
                                    pygame.time.delay(500)
                                    
                                    # 电脑落子
                                    comp_move = self.get_computer_move()
                                    if comp_move:
                                        self.make_move(comp_move[0], comp_move[1], COMPUTER)
                                        self.current_player = PLAYER
                                        self.draw_board()
            
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = GomokuGame()
    game.run()